package com.md.config.database.note;

import com.baomidou.mybatisplus.core.MybatisParameterHandler;
import com.google.common.collect.Maps;
import com.md.util.ClassUtils;
import com.md.util.ObjectUtil;
import lombok.extern.slf4j.Slf4j;
import org.apache.ibatis.executor.parameter.ParameterHandler;
import org.apache.ibatis.executor.statement.StatementHandler;
import org.apache.ibatis.mapping.BoundSql;
import org.apache.ibatis.mapping.MappedStatement;
import org.apache.ibatis.plugin.Interceptor;
import org.apache.ibatis.plugin.Intercepts;
import org.apache.ibatis.plugin.Invocation;
import org.apache.ibatis.plugin.Signature;
import org.apache.ibatis.scripting.defaults.DefaultParameterHandler;
import org.springframework.util.ReflectionUtils;

import java.lang.reflect.Field;
import java.sql.Connection;
import java.util.Map;
import java.util.Optional;
import java.util.concurrent.Callable;

/**
 * @author 刘骄阳
 * 2023-12-20 15:41
 */
@Slf4j
@Intercepts({
  /*
   * SQL 准备阶段调用的。MyBatis 会在执行 SQL 之前，先通过 StatementHandler.prepare() 方法准备 SQL 语句。
   */
  @Signature(
    type = StatementHandler.class,
    method = "prepare",
    args = {Connection.class, Integer.class}
  )})
public class NoteInterceptor implements Interceptor {

  private final NoteSQLConfig noteSQLConfig;
  private final String noteSqlFormat = "/* [***SQL NOTE***]  %s */";
  private final String noteSqlContentFormat = " %s: %s ";

  private final Map<String, Callable<String>> noteCallableMap;

  public NoteInterceptor(final NoteSQLConfig noteSQLConfig) {
    noteCallableMap = Maps.newHashMap();
    noteCallableMap.put("threadName", () -> Thread.currentThread().getName());
    this.noteSQLConfig = noteSQLConfig;
  }


  public void appendNoteMap(final String key, final Callable<String> task) {
    noteCallableMap.put(key, task);
  }

  public void remark(final BoundSql boundSql) {
    log.info("exec SQL: {}", boundSql.getSql());
  }

  public void note(final MappedStatement ms, final BoundSql boundSql) {
    try {
      Field field = ClassUtils.getField("sql", boundSql);
      if (ObjectUtil.isEmpty(field)) {
        return;
      }
      StringBuilder sb = new StringBuilder();
      Optional.ofNullable(ms).ifPresent(mappedStatement -> {
        sb.append(String.format(noteSqlContentFormat, "id", ms.getId()));
      });
      noteCallableMap.forEach((key, valueTask) -> {
        try {
          sb.append(String.format(noteSqlContentFormat, key, valueTask.call()));
        } catch (Exception e) {
          log.debug("拼接失败: {}", e.getMessage(), e);
        }
      });
      Optional.ofNullable(noteSQLConfig).ifPresent(bean -> bean.append(sb));
      String sql = boundSql.getSql();
      ReflectionUtils.makeAccessible(field);
      ReflectionUtils.setField(field, boundSql, sql + String.format(noteSqlFormat, sb));
    } catch (Exception e) {
      log.error("签名失败: {}", e.getMessage(), e);
    }
  }

  @Override
  public Object intercept(final Invocation invocation) throws Exception {
    // StatementHandler
    final StatementHandler sh = (StatementHandler) invocation.getTarget();
    BoundSql boundSql = sh.getBoundSql();
    ParameterHandler parameterHandler = sh.getParameterHandler();
    MappedStatement ms;
    if (parameterHandler instanceof MybatisParameterHandler) {
      ms = getMapperStatement((MybatisParameterHandler) parameterHandler);
    } else if (parameterHandler instanceof DefaultParameterHandler) {
      ms = getMapperStatement((DefaultParameterHandler) parameterHandler);
    } else {
      ms = null;
    }
    remark(boundSql);
    note(ms, boundSql);
    return invocation.proceed();
  }

  public MappedStatement getMapperStatement(final MybatisParameterHandler handler) {
    return (MappedStatement) ClassUtils.getFieldValue("mappedStatement", handler);
  }

  public MappedStatement getMapperStatement(final DefaultParameterHandler handler) {
    return (MappedStatement) ClassUtils.getFieldValue("mappedStatement", handler);
  }
}
